import random
import math
from typing import List, Tuple

def assign_cluster(point: List[float], centers: List[List[float]]) -> int:
    #将数据点分配到最近的聚类中心
    return min(range(len(centers)),
               key=lambda i: sum((p - c) ** 2 for p, c in zip(point, centers[i])))

def kmeans(data: List[List[float]], k: int, eps: float = 1e-4, max_iter: int = 100) -> Tuple[
    List[List[float]], List[int]]:
    #K-means聚类算法实现
    if len(data) < k:
        raise ValueError("数据点数量少于聚类数k")

    centers = random.sample(data, k)
    assignments = []

    for _ in range(max_iter):
        # 分配聚类
        assignments = [assign_cluster(p, centers) for p in data]

        # 计算新中心
        new_centers = []
        for i in range(k):
            points = [data[j] for j, c in enumerate(assignments) if c == i]
            if not points:
                new_centers.append(random.choice(data))
                continue
            new_centers.append([sum(d) / len(points) for d in zip(*points)])

        # 检查收敛
        if max(math.sqrt(sum((o - n) ** 2 for o, n in zip(old, new)))
               for old, new in zip(centers, new_centers)) < eps:
            break
        centers = new_centers

    return centers, assignments


def calculate_sse(data: List[List[float]], centers: List[List[float]], assignments: List[int]) -> float:
    # 计算误差平方和
    return sum(sum((p - c) ** 2 for p, c in zip(data[i], centers[assignments[i]]))
               for i in range(len(data)))


# 测试
if __name__ == "__main__":
    random.seed(66)
    data = []
    # 生成3类测试数据
    for mu in [(1, 1), (2, 2), (3, 3)]:
        data.extend([[random.gauss(m, 0.3) for m in mu] for _ in range(30)])

    centers, assignments = kmeans(data, 3)

    print("聚类中心:", [list(map(round, c, [2] * 2)) for c in centers])
    print("各聚类数量:", [assignments.count(i) for i in range(3)])
    print("SSE:", f"{calculate_sse(data, centers, assignments):.4f}")